package lock.etcd;

import io.etcd.jetcd.*;
import io.etcd.jetcd.kv.TxnResponse;
import io.etcd.jetcd.lease.LeaseKeepAliveResponse;
import io.etcd.jetcd.op.Cmp;
import io.etcd.jetcd.op.CmpTarget;
import io.etcd.jetcd.op.Op;
import io.etcd.jetcd.options.DeleteOption;
import io.etcd.jetcd.options.PutOption;
import io.etcd.jetcd.watch.WatchEvent;
import io.grpc.stub.StreamObserver;

import java.nio.charset.StandardCharsets;
import java.util.UUID;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicBoolean;

public class EtcdLock {

    private String key;
    private Client client;
    private volatile long leaseId = 0L;
    private String holderId;
    public EtcdLock(Client client, String lockName) {
        this.client = client;
        this.holderId = getHolderId();
        this.key = "/lock/" + lockName;
    }

    public void lock() throws ExecutionException, InterruptedException {
        Lease leaseClient = client.getLeaseClient();
        leaseId = leaseClient.grant(10).join().getID();

        leaseClient.keepAlive(leaseId, new StreamObserver<>() {
            @Override
            public void onNext(LeaseKeepAliveResponse leaseKeepAliveResponse) {

            }

            @Override
            public void onError(Throwable throwable) {

            }

            @Override
            public void onCompleted() {

            }
        });
        ByteSequence keyBs = ByteSequence.from(key, StandardCharsets.UTF_8);
        ByteSequence valueBs = ByteSequence.from(holderId, StandardCharsets.UTF_8);
        PutOption putOpt = PutOption.newBuilder().withLeaseId(leaseId).build();

        while (true) {

            // 原子性CAS key不存在则写入成功
            Cmp cmp = new Cmp(keyBs, Cmp.Op.EQUAL, CmpTarget.version(0));
            Op put = Op.put(keyBs, valueBs, putOpt);
            Txn txn = client.getKVClient().txn();
            TxnResponse resp = txn.If(cmp).Then(put).commit().get();

            if (resp.isSucceeded()) {
                return;
            }

            CountDownLatch latch = new CountDownLatch(1);
            AtomicBoolean triggered  =  new AtomicBoolean(false);

            Watch.Watcher watcher = client.getWatchClient().watch(keyBs, watchResponse -> {
                for (WatchEvent event : watchResponse.getEvents()) {
                    if (event.getEventType() == WatchEvent.EventType.DELETE) {
                        if (triggered.compareAndSet(false, true)) {
                            latch.countDown();
                        }
                    }
                }
            });


            latch.await();
            // 取消监听
            watcher.close();
        }

    }

    public void unlock() {
        if (leaseId == 0L) {
            return;
        }
        ByteSequence keyBs = ByteSequence.from(key, StandardCharsets.UTF_8);
        ByteSequence valueBs = ByteSequence.from(holderId, StandardCharsets.UTF_8);

        Cmp cmp = new Cmp(keyBs, Cmp.Op.EQUAL, CmpTarget.value(valueBs));
        Op delete = Op.delete(keyBs, DeleteOption.DEFAULT);
        Txn txn = client.getKVClient().txn();
        txn.If(cmp).Then(delete).commit();

        long _leaseId = leaseId;
        leaseId = 0L;
        client.getLeaseClient().revoke(_leaseId);
    }


    public String getHolderId() {
        return Thread.currentThread().threadId() +
                UUID.randomUUID().toString();
    }

    public static void main(String[] args) {
        Client client = Client.builder().endpoints("http://localhost:2379").build();
        EtcdLock lock = new EtcdLock(client, "my-lock");
    }
}
